from typing import List
import numpy as np
import manim as mnm

from investment.time_series import IndexBasketTimeSeries, IndexTimeSeries, find_in_series
from investment.constants import A_INDEX_DICT
from investment.date_utilities import get_date_by_month, get_date_by_week, segment_week_by_month, get_day_interval
from investment.operation import calc_reward_rate
from misc.emotion_utils import calc_emotion_series
from misc.character_design import SmallYan, Emo, LookAt, Yan, PiCreatureSays, RemovePiCreatureBubble
from visual.common_scenes import TimeSeriesAxes, MY_FONT


RewardToEmotionDict = {"low": -0.1, "up": 0.2}


def always_redraw(func):
    mob = func()
    mob.add_updater(lambda m: mob.become(func()))
    return mob


class EmotionAxes(TimeSeriesAxes):
    def __init__(self, value_arr, dates, title, height):
        super().__init__(value_arr, title, height)
        self.dates = dates
        self.character = None

    def get_line_to_x_axis(self):
        return self.axes.get_vertical_line(self.dot.get_center(), line_config={"dashed_ratio": 0.85})

    def get_x_label(self):
        axis = self.axes.get_axis(0)
        axis_projection = axis.get_projection(self.dot.get_center())
        date = self.dates[int(self.axes.point_to_coords(self.dot.get_center())[0])]
        return mnm.Tex(date, font_size=18).next_to(axis_projection, direction=mnm.DOWN)

    def set_dot(self, date_idx, color=mnm.DARK_BLUE):
        self.dot = mnm.Dot(self.index_to_point(date_idx), color=color)
        self.trace = mnm.TracedPath(self.dot.get_center, dissipating_time=0.3, stroke_opacity=[0, 1])

    def set_character(self):
        assert self.dot is not None
        self.character = SmallYan().next_to(self.dot, direction=mnm.RIGHT)
        self.character.add_updater(lambda mobject: mobject.next_to(self.dot, direction=mnm.RIGHT))


class EmotionAxesScene(mnm.MovingCameraScene):
    def setup(self):
        self.index_boxes = IndexBasketTimeSeries("2019-10-01", "2021-10-01",
                                                 [A_INDEX_DICT['中证500'], A_INDEX_DICT['沪深300']])
        self.net_value = self.index_boxes.get_merged_attribute('get_net_value')
        self.month_record_day = 7
        self.week_record_day = 4
        axes_height = mnm.config.frame_height // 3
        self.axes_month = EmotionAxes(self.net_value, self.index_boxes.dates_series,
                                      r"沪深300:中证500 = 1:1 ; 月记账", axes_height)
        month_indexes, month_date = get_date_by_month(self.index_boxes.dates_series, 1)
        self.axes_month.add_date_ticks(dict(zip(month_indexes, month_date)))
        self.axes_week = EmotionAxes(self.net_value, self.index_boxes.dates_series,
                                     r"沪深300:中证500 = 1:1 ; 周记账", axes_height)
        self.axes_day = EmotionAxes(self.net_value, self.index_boxes.dates_series,
                                    r"沪深300:中证500 = 1:1 ; 日记账", axes_height)

    def construct(self):
        self.add(self.axes_month.axes_group)
        self.mov_net_value_by_month()
        self.play(self.axes_month.axes_group.animate.shift(2.5 * mnm.UP))
        self.add(self.axes_week.axes_group)
        self.mov_net_value_by_week()
        self.play(*[mnm.FadeOut(self.axes_month.axes_group), self.axes_week.axes_group.animate.shift(2.5 * mnm.UP)])
        self.add(self.axes_day.axes_group)
        self.mov_net_value_by_day()
        self.play(*[mnm.FadeIn(self.axes_month.axes_group, self.axes_week.axes_group.animate.shift(2.5 * mnm.DOWN),
                               self.axes_day.axes_group.shift(2.5 * mnm.DOWN))])
        # self.add_table()
        self.wait(0.5)

    def mov_net_value_by_month(self):
        month_indexes, month_date = get_date_by_month(self.index_boxes.dates_series, self.month_record_day)
        self.axes_month.set_dot(month_indexes[0])
        self.axes_month.set_character()
        self.play(mnm.FadeIn(self.axes_month.dot))
        self.add(self.axes_month.traced, self.axes_month.character)
        line_x_axis = always_redraw(self.axes_month.get_line_to_x_axis)
        x_label = always_redraw(self.axes_month.get_x_label)
        emotion = calc_emotion_series(self.net_value[month_indexes])
        self.add(line_x_axis, x_label)

        self.month_dots = mnm.VGroup()
        self.month_lines = mnm.VGroup()
        for idx, series_idx in enumerate(month_indexes[1:]):
            pt = self.axes_month.index_to_point(series_idx)
            dot_cp = self.axes_month.dot.copy()
            self.add(dot_cp)
            self.play(mnm.ApplyMethod(self.axes_month.dot.move_to, pt, run_time=0.5))
            self.play(Emo(self.axes_month.character, emotion[idx + 1]), run_time=0.5)
            self.play(LookAt(self.axes_month.character, self.axes_month.dot), run_time=0.1)
            self.play(mnm.FadeOut(dot_cp), run_time=0.1)
            if idx > 0:
                line = mnm.Line(dot_cp.get_center(), self.axes_month.dot.get_center())
                self.month_lines.add(line)
            self.month_dots.add(dot_cp)
        self.month_dots.add(self.axes_month.dot)
        self.remove(self.axes_month.dot, self.axes_month.character)

        # Show all lines
        self.play(mnm.Create(self.month_dots))
        self.play(mnm.Create(self.month_lines))
        self.axes_month.axes_group.add(*[self.month_dots, self.month_lines])

    def mov_net_value_by_week(self):
        # Bucket filling
        month_indexes, _ = get_date_by_month(self.index_boxes.dates_series, self.month_record_day)
        week_indexes, _ = get_date_by_week(self.index_boxes.dates_series, self.week_record_day)
        month_week_bucket = segment_week_by_month(month_indexes, week_indexes)

        self.axes_month.set_dot(week_indexes[0])
        dot_month_end = self.axes_month.dot
        self.axes_week.set_dot(week_indexes[0])
        self.axes_week.set_character()
        self.play(*[mnm.FadeIn(self.axes_month.dot), mnm.FadeIn(self.axes_week.dot)])
        line_x_axis = always_redraw(self.axes_week.get_line_to_x_axis)
        x_label = always_redraw(self.axes_week.get_x_label)
        self.add(line_x_axis, x_label)
        emotion = calc_emotion_series(self.net_value[week_indexes])

        self.week_dots = mnm.VGroup()
        self.week_lines = mnm.VGroup()
        idx = 0
        for month_idx_end in month_week_bucket.keys():
            self.play(*[self.axes_month.dot.animate.move_to(dot_month_end),
                        dot_month_end.animate.move_to(self.axes_month.index_to_point(month_idx_end))])
            dashed_line_month = mnm.DashedLine(self.axes_month.dot, dot_month_end)
            self.play(mnm.Create(dashed_line_month))
            for week_idx in month_week_bucket[month_idx_end]:
                pt = self.axes_week.index_to_point(week_idx)
                dot_cp = self.axes_week.dot.copy()
                self.add(dot_cp)
                self.play(mnm.ApplyMethod(self.axes_week.dot.move_to, pt, run_time=0.3))
                self.play(Emo(self.axes_week.character, emotion[idx + 1]), run_time=0.3)
                self.play(LookAt(self.axes_week.character, self.axes_week.dot), run_time=0.05)
                self.play(mnm.FadeOut(dot_cp), run_time=0.05)
                if idx > 0:
                    line = mnm.Line(dot_cp.get_center(), self.axes_week.dot.get_center())
                    self.week_lines.add(line)
                self.week_dots.add(dot_cp)
                idx += 1
            self.remove(dashed_line_month)
        self.week_dots.add(self.axes_week.dot)
        self.remove(self.axes_week.dot, self.axes_week.character)

        # Show all lines
        self.play(mnm.Create(self.week_dots))
        self.play(mnm.Create(self.week_lines))
        self.axes_week.axes_group.add(*[self.week_dots, self.week_lines])

    def mov_net_value_by_day(self):
        # Bucket filling
        month_indexes, _ = get_date_by_month(self.index_boxes.dates_series, self.month_record_day)

        self.axes_month.set_dot(month_indexes[0])
        dot_month_end = self.axes_month.dot.copy()
        self.axes_day.set_dot(month_indexes[0])
        self.axes_day.set_character()
        self.play(*[mnm.FadeIn(self.axes_month.dot), mnm.FadeIn(self.axes_day.dot)])
        line_x_axis = always_redraw(self.axes_day.get_line_to_x_axis)
        x_label = always_redraw(self.axes_day.get_x_label)
        self.add(line_x_axis, x_label)
        emotion = calc_emotion_series(self.net_value[month_indexes[0]:month_indexes[-1]])

        # camera save state
        self.camera.frame.save_state()
        def camera_move_to(mob):
            mob.move_to(self.axes_day.dot)
        self.play(self.camera.frame.animate.scale(0.4).move_to(self.axes_day.dot))
        self.camera.frame.add_updater(camera_move_to)

        self.day_dots = mnm.VGroup()
        self.day_lines = mnm.VGroup()
        for idx, month_idx_end in enumerate(month_indexes):
            self.play(*[self.axes_month.dot.animate.move_to(dot_month_end),
                        dot_month_end.animate.move_to(self.axes_month.index_to_point(month_idx_end))])
            dashed_line_month = mnm.DashedLine(self.axes_month.dot, dot_month_end)
            self.play(mnm.Create(dashed_line_month))
            for day_idx in range(month_indexes[idx-1], month_indexes[idx]):
                pt = self.axes_day.index_to_point(day_idx)
                dot_cp = self.axes_day.dot.copy()
                self.add(dot_cp)
                self.play(mnm.ApplyMethod(self.axes_day.dot.move_to, pt, run_time=0.3))
                self.play(Emo(self.axes_day.character, emotion[idx + 1]), run_time=0.3)
                self.play(LookAt(self.axes_day.character, self.axes_day.dot), run_time=0.05)
                self.play(mnm.FadeOut(dot_cp), run_time=0.05)
                if idx > 0:
                    line = mnm.Line(dot_cp.get_center(), self.axes_day.dot.get_center())
                    self.day_lines.add(line)
                self.day_dots.add(dot_cp)
                idx += 1
            self.remove(dashed_line_month)
        self.day_dots.add(self.axes_day.dot)
        self.remove(self.axes_day.dot, self.axes_day.character)
        self.camera.frame.remove_updater(camera_move_to)
        self.play(mnm.Restore(self.camera.frame))
        # Show all lines
        self.play(mnm.Create(self.day_dots))
        self.play(mnm.Create(self.day_lines))
        self.axes_day.axes_group.add(*[self.day_dots, self.day_lines])

    def add_table(self, axes, emotion: List):
        pi = Yan(height=1)
        table = mnm.Table(
            [["×3"], ["×1"], ["×5"]],
            row_labels=[pi, pi.copy(), pi.copy()],
            col_labels=[mnm.Text("Times")])
        ent = table.get_entries()
        for item in ent:
            item.set_color(mnm.random_bright_color())
        self.play(mnm.Create(table))


def reward_to_emotion(reward: float):
    res = "心情: "
    if reward < RewardToEmotionDict["low"]:
        res += "Sad."
    elif reward > RewardToEmotionDict["up"]:
        res += "Excited!"
    else:
        res += "Peaceful."
    return res


def ch(num):
    chinese = ["一", "二", "三", "四", "五", "六", "七", "八", "九", "十"]
    return chinese[num]


class EmotionSeriesSegments:
    def __init__(self, series, start_pe, height=mnm.config.frame_height // 2):
        assert isinstance(series, IndexTimeSeries)
        self.series = series
        self.height = height
        pe_arr = series.get_fundamental_metrics("PE")
        self.series_length = len(pe_arr)
        self.spacer = 100
        # note: from start to end, add final EOF length.
        self.start_indexes = np.append(find_in_series(pe_arr, start_pe, filter_pad=self.spacer), self.series_length)
        self.seg_count = len(self.start_indexes) - 1

    def __iter__(self):
        self.iterator = iter(
            map(lambda i: AxesSegment(self.series, self.start_indexes[i], self.start_indexes[i + 1],
                                      f"时间{ch(i)}: {self.series.dates_series[self.start_indexes[i]]}",
                                      self.height), range(self.seg_count)))
        return self.iterator


class AxesSegment(EmotionAxes):
    def __init__(self, series, st, ed, title, height):
        assert isinstance(series, IndexTimeSeries)
        value_arr = series.get_fundamental_metrics("CP")[st:ed]
        dates = series.dates_series[st:ed]
        length = len(value_arr)
        top = np.nanargmax(value_arr)
        margin = 30
        self.ed = top + margin if top + margin < length else length - 1
        # find lowest price or payback
        low = np.nanargmin(value_arr[:top])
        if low > margin:
            low = low
        else:
            holding_zone = 1.01
            # give me the money, back to home.
            low = np.nanargmax(value_arr[:top] > value_arr[0]*holding_zone)
        self.journey_nodes = {"low": low, "top": top, "end": self.ed}
        super().__init__(value_arr[:self.ed+1], dates[:self.ed+1], title=title, height=height)

        self.dot = mnm.Dot(self.index_to_point(0), color=mnm.BLUE)
        self.trace = mnm.TracedPath(self.dot.get_center)
        self.character = SmallYan().next_to(self.dot, direction=mnm.RIGHT)
        self.character.add_updater(lambda mobject: mobject.next_to(self.dot, direction=mnm.RIGHT))

        line_x_axis = always_redraw(self.get_line_to_x_axis)
        x_label = always_redraw(self.get_x_label)
        self.redraw_x = mnm.Group(line_x_axis, x_label)
        self.duration_text = None

    def animate_character_emotion(self, last_position, position):
        reward_rate = calc_reward_rate(self.value_series[0], self.value_series[position])
        text = f"记账:{reward_rate:+.1%}\n"
        """
        回报率非线性规律:
        a = (ed0-st)/st
        b = (ed1-ed0)/ed0
        c = (ed1-st)/st
        c/a = (ed1-st)/(ed0-st)
        a/c = ed0-st/ed1-st
        c-a = (ed1-ed0)/st
        so:
        1/(c-a) - 1/b = (st-ed0)/ed1-ed0 = a/c
        """
        lost_rate = reward_rate
        if last_position != 0:
            lost_rate = reward_rate - calc_reward_rate(self.value_series[0], self.value_series[last_position])
            text += f"比上次:{lost_rate:+.1%}\n"
        text += reward_to_emotion(lost_rate)
        saying = PiCreatureSays(self.character, text, bubble_kwargs={'height': 1, 'width': 2})
        self.character.generate_target()
        self.character.target.look_at(self.index_to_point(last_position))
        # emo should be later by saying, conver to -1~1 by lost rate.
        self.character.target.emo(min(1, max(-1, lost_rate / RewardToEmotionDict['up'])))
        change_mode = mnm.MoveToTarget(self.character)
        return mnm.AnimationGroup(saying, change_mode)

    def remove_bubble(self):
        return RemovePiCreatureBubble(self.character)

    def animate_moving(self, st, ed, color=mnm.WHITE):
        moving_line = self.generate_wandering_line(st, ed, color)
        return mnm.MoveAlongPath(self.dot, moving_line)

    def get_duration_text(self):
        interval = get_day_interval(self.dates[0], self.dates[self.ed])
        self.duration_text = mnm.Text(f"持有{interval}天", font_size=18, font=MY_FONT).next_to(
            self.axes_group.submobjects[1], direction=mnm.DOWN, buff=mnm.MED_LARGE_BUFF)
        return self.duration_text

    def cache_position(self):
        dot = self.dot.copy()
        line = self.get_line_to_x_axis()
        date = self.get_x_label()
        return mnm.VGroup(dot, line, date)


class HoldingSceneTest(mnm.Scene):
    def construct(self):
        a_market = IndexTimeSeries(A_INDEX_DICT['沪深300'], "2002-07-01", "2022-07-01", fill_flag=True)
        title = mnm.Text("看看相对低估时(PE≈10)买入:", font_size=18, font=MY_FONT)
        title.scale(1.5)
        title.to_edge(mnm.UL, buff=mnm.MED_LARGE_BUFF)
        self.add(title)
        emotion_axes = EmotionSeriesSegments(a_market, 12)
        for axis in iter(emotion_axes):
            self.add(axis.axes_group, axis.trace, axis.character, axis.redraw_x)
            last_pos = 0
            for key, pos in axis.journey_nodes.items():
                self.play(axis.animate_moving(last_pos, pos))
                self.play(axis.animate_character_emotion(last_pos, pos))
                self.play(axis.remove_bubble())
                last_pos = pos
            self.play(mnm.Create(axis.get_duration_text(), remover=True))
            self.wait(1)
            self.remove(axis.axes_group, axis.trace, axis.dot, axis.character, axis.redraw_x)
            break
        conclusion = "向山顶攀登的过程中,会有种种经历和感受. 耐心持有,静待花开."
        self.play(mnm.Write(mnm.Text(conclusion, font_size=18, font=MY_FONT).next_to(
            title, direction=mnm.DOWN, buff=mnm.LARGE_BUFF)))


if __name__ == "__main__":
    with mnm.tempconfig({"quality": "low_quality", "preview": True}):
        scene = HoldingSceneTest()
        scene.render()
